#!/bin/bash

# ====================================================================================
# CSE_KGAN: LLM驱动的语义与结构双重增强知识图谱推荐系统 - 自动化工作流脚本 (v5.6)
# ====================================================================================

# --- 颜色和样式定义 (省略) ---
RED='\033[1;31m'
GREEN='\033[1;32m'
YELLOW='\033[1;33m'
BLUE='\033[1;34m'
PURPLE='\033[1;35m'
CYAN='\033[1;36m'
NC='\033[0m'
BOLD='\033[1m'
print_header() {
	echo -e "\n${PURPLE}${BOLD}==========================================================${NC}"
	echo -e "${PURPLE}${BOLD}$1${NC}"
	echo -e "${PURPLE}${BOLD}==========================================================${NC}"
}
print_section() { echo -e "\n${CYAN}${BOLD}=== $1 ===${NC}"; }
print_success() { echo -e "${GREEN}✅ SUCCESS: $1${NC}"; }
print_warning() { echo -e "${YELLOW}⚠️  WARNING: $1${NC}"; }
print_error() { echo -e "${RED}❌ ERROR: $1${NC}"; }
print_param() { printf "  ${CYAN}%-22s${NC} : ${GREEN}%s${NC}\n" "$1" "$2"; }
print_progress() { echo -e "\n${BLUE}${BOLD}🚀 $1...${NC}"; }

# --- 帮助信息函数 (省略) ---

# --- 1. 默认与命令行参数配置 ---
# ... (省略，与上一版相同) ...
DATASET_NAME="amazon-book"
GPU_ID=0
EMBED_DIM=64
CONV_DIMS="[64, 32, 16]"
MESS_DROPOUT="[0.1, 0.1, 0.1]"
LEARNING_RATE=1e-4
CF_BATCH_SIZE=4096
N_EPOCH=500
EVALUATE_EVERY=5
STOPPING_STEPS=10
CF_L2LOSS_LAMBDA=1e-5
CONTRASTIVE_LAMBDA=0.1
SEED=2024
VALID_METRIC="Recall@20"
Ks="[20, 40]"
LLM_SAMPLE_SIZE=200
while [[ $# -gt 0 ]]; do
	case $1 in -d | --dataset)
		DATASET_NAME="$2"
		shift 2
		;;
	-g | --gpu_id)
		GPU_ID="$2"
		shift 2
		;;
	*)
		print_error "未知参数: $1"
		exit 1
		;;
	esac
done

# --- 2. 环境变量与目录设置 ---
print_progress "初始化环境与目录"
if [ -f ".env" ]; then export $(grep -v '^#' .env | xargs); fi
export HF_ENDPOINT=https://hf-mirror.com
export CUDA_VISIBLE_DEVICES=$GPU_ID
print_success "Hugging Face 镜像已设置为: ${HF_ENDPOINT}"
SAVE_DIR="saved/${DATASET_NAME}/embed-dim${EMBED_DIM}_lr${LEARNING_RATE}_cl${CONTRASTIVE_LAMBDA}/"
LOG_DIR="logs/"
mkdir -p ${SAVE_DIR} ${LOG_DIR}
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOG_FILE="${LOG_DIR}/cse-kgan_${DATASET_NAME}_${TIMESTAMP}.log"
print_success "目录与日志文件设置完毕"

# --- 3. 环境检查 ---
print_section "环境检查"
if [ ! -d "data/${DATASET_NAME}" ]; then
	print_warning "数据集目录 data/${DATASET_NAME} 不存在"
else
	print_success "数据集目录检查通过"
fi
PYTHON_VERSION=$(python --version 2>&1)
print_success "Python 环境: ${PYTHON_VERSION}"

# --- 自动化LLM数据预处理工作流 ---
print_header "自动化 LLM 数据预处理检查"
USER_PROFILE_PT="data/${DATASET_NAME}/user_profiles.pt"
ITEM_PROFILE_PT="data/${DATASET_NAME}/item_profiles.pt"
USER_PROFILE_JSON="data/${DATASET_NAME}/user_profiles.json"
ITEM_PROFILE_JSON="data/${DATASET_NAME}/item_profiles.json"
CS_KG_FILE="data/${DATASET_NAME}/llm_enhanced_kg.txt"

# 【关键修正】: 定义一个更健壮的命令执行函数
run_command() {
	print_progress "$2"
	eval $1
	if [ $? -ne 0 ]; then
		print_error "$2 失败，脚本终止。请检查上面的错误日志。"
		exit 1
	fi
}

if [ ! -s "${USER_PROFILE_PT}" ] || [ ! -s "${ITEM_PROFILE_PT}" ]; then
	print_warning "嵌入文件 (.pt) 未找到或为空。启动自动生成流程..."
	if [ ! -s "${USER_PROFILE_JSON}" ] || [ ! -s "${ITEM_PROFILE_JSON}" ]; then
		run_command "python llm_enhancement/1_generate_profiles.py --data_name ${DATASET_NAME} --sample_size ${LLM_SAMPLE_SIZE}" "步骤 1/3: 生成文本画像"
	else
		print_success "步骤 1/3: 文本画像 (.json) 已存在且非空。"
	fi
	run_command "python llm_enhancement/3_encode_profiles.py --data_name ${DATASET_NAME} --gpu_id ${GPU_ID}" "步骤 2/3: 编码画像文件"
	print_success "嵌入文件生成完毕。"
else
	print_success "语义画像嵌入文件 (.pt) 已存在且非空，跳过生成。"
fi

if [ ! -s "${CS_KG_FILE}" ]; then
	print_warning "常识知识图谱 (${CS_KG_FILE}) 未找到或为空。"
	run_command "python llm_enhancement/2_generate_cs_kg.py --data_name ${DATASET_NAME} --sample_size ${LLM_SAMPLE_SIZE}" "步骤 3/3: 生成常识知识图谱"
	print_success "常识知识图谱生成完毕。"
else
	print_success "常识知识图谱文件已存在且非空，跳过生成。"
fi

# --- 6. 打印最终配置 ---
print_header "CSE-KGAN 训练启动配置"
print_section "核心参数"
print_param "数据集" "${DATASET_NAME}"
print_param "模型保存目录" "${SAVE_DIR}"
print_param "日志文件" "${LOG_FILE}"
print_section "模型与训练参数"
print_param "嵌入维度" "${EMBED_DIM}"
print_param "GNN层维度" "${CONV_DIMS}"
print_param "学习率" "${LEARNING_RATE}"
print_param "批次大小" "${CF_BATCH_SIZE}"
print_param "对比损失权重" "${CONTRASTIVE_LAMBDA}"
print_param "早停指标" "${VALID_METRIC}"
print_param "最大轮数" "${N_EPOCH}"

# --- 7. 执行主训练流程 ---
print_header "启动主训练流程"
print_progress "开始后台训练进程"
nohup python -u main.py \
	--data_name "${DATASET_NAME}" \
	--gpu_id "${GPU_ID}" \
	--embed_dim "${EMBED_DIM}" \
	--conv_dim_list "${CONV_DIMS}" \
	--mess_dropout "${MESS_DROPOUT}" \
	--lr "${LEARNING_RATE}" \
	--cf_batch_size "${CF_BATCH_SIZE}" \
	--n_epoch "${N_EPOCH}" \
	--evaluate_every "${EVALUATE_EVERY}" \
	--stopping_steps "${STOPPING_STEPS}" \
	--cf_l2loss_lambda "${CF_L2LOSS_LAMBDA}" \
	--contrastive_lambda "${CONTRASTIVE_LAMBDA}" \
	--seed "${SEED}" \
	--save_dir "${SAVE_DIR}" \
	--valid_metric "${VALID_METRIC}" \
	--Ks "${Ks}" >"${LOG_FILE}" 2>&1 &
PID=$!
echo ""
print_success "训练进程已在后台启动!"
print_param "进程PID" "${PID}"
print_header "训练监控指南"
echo -e "${GREEN}📊 实时查看日志:${NC}"
echo -e "  ${CYAN}tail -f ${LOG_FILE}${NC}"
echo -e "\n${YELLOW}🔧 系统监控命令:${NC}"
echo -e "  ${CYAN}nvidia-smi${NC}"
echo -e "  ${CYAN}ps -p ${PID} -o pid,pcpu,pmem,etime${NC}"
echo -e "\n${RED}🛑 停止训练命令:${NC}"
echo -e "  ${CYAN}kill ${PID}${NC}"
echo -e "\n${BLUE}📈 关键指标查看:${NC}"
echo -e "  ${CYAN}grep -E 'Evaluation|Best|Epoch' ${LOG_FILE} | tail -n 10${NC}"
echo -e "  ${CYAN}grep -i 'error\\|warning\\|nan' ${LOG_FILE}${NC}"
print_header ""
sleep 1
if ! ps -p $PID >/dev/null; then
	print_error "训练进程启动失败，请立即检查日志文件: ${LOG_FILE}"
	exit 1
fi
